#include "lua.h"
#include "lauxlib.h"
#include "cm_uart.h"
#include "cm_iomux.h"
#include "lua_test_entry.h"
#include <string.h>
#include <stdlib.h>

#define MAX_UARTS CM_UART_DEV_NUM

// Structure to hold callback info for each UART
typedef struct {
    int on_receive_ref;
} lua_uart_cb_t;

lua_uart_cb_t g_uart_cbs[MAX_UARTS];

// --- C Callback for UART events ---
static void uart_event_callback(void *param, uint32_t evt_id)
{
    cm_uart_dev_e dev = (cm_uart_dev_e)param;
    if (dev >= MAX_UARTS) return;

    if (evt_id & CM_UART_EVENT_TYPE_RX_ARRIVED)
    {
        lua_event_t event;
        event.id = LUA_EVENT_UART_RX;
        event.data.ival = dev; // Pass the UART device number
        lua_event_post(&event);
    }
}

// Lua: uart.setup(id, baudrate, databits, parity, stopbits)
static int l_uart_setup(lua_State *L)
{
    cm_uart_dev_e dev = (cm_uart_dev_e)luaL_checkinteger(L, 1);

    // --- STEP 1: Configure Pin Muxing (THE CRITICAL STEP) ---
    if (dev == CM_UART_DEV_0) {
        cm_iomux_set_pin_func(CM_IOMUX_PIN_17, CM_IOMUX_FUNC_FUNCTION1);
        cm_iomux_set_pin_func(CM_IOMUX_PIN_18, CM_IOMUX_FUNC_FUNCTION1);
    } else if (dev == CM_UART_DEV_1) {
        cm_iomux_set_pin_func(CM_IOMUX_PIN_28, CM_IOMUX_FUNC_FUNCTION1);
        cm_iomux_set_pin_func(CM_IOMUX_PIN_29, CM_IOMUX_FUNC_FUNCTION1);
    } else if (dev == CM_UART_DEV_2) {
        cm_iomux_set_pin_func(CM_IOMUX_PIN_38, CM_IOMUX_FUNC_FUNCTION2);
        cm_iomux_set_pin_func(CM_IOMUX_PIN_39, CM_IOMUX_FUNC_FUNCTION2);
    }
    // Add configurations for other UARTs here if needed

    // --- STEP 2: Configure UART parameters ---
    cm_uart_cfg_t cfg;
    cfg.baudrate = luaL_checkinteger(L, 2);
    cfg.byte_size = (cm_uart_byte_size_e)luaL_optinteger(L, 3, CM_UART_BYTE_SIZE_8);
    cfg.parity = (cm_uart_parity_e)luaL_optinteger(L, 4, CM_UART_PARITY_NONE);
    cfg.stop_bit = (cm_uart_stop_bit_e)luaL_optinteger(L, 5, CM_UART_STOP_BIT_ONE);
    cfg.flow_ctrl = CM_UART_FLOW_CTRL_NONE;
    cfg.is_lpuart = 0;
    cfg.rxrb_buf_size = 0; // Use default

    // Register event before opening
    cm_uart_event_t uart_event = {
        .event_type = CM_UART_EVENT_TYPE_RX_ARRIVED | CM_UART_EVENT_TYPE_RX_OVERFLOW,
        .event_param = (void*)dev,
        .event_entry = uart_event_callback,
    };
    cm_uart_register_event(dev, &uart_event);

    int ret = cm_uart_open(dev, &cfg);
    lua_pushboolean(L, ret == 0);
    return 1;
}

// Lua: uart.write(id, data)
static int l_uart_write(lua_State *L)
{
    cm_uart_dev_e dev = (cm_uart_dev_e)luaL_checkinteger(L, 1);
    size_t len;
    const char *data = luaL_checklstring(L, 2, &len);

    int32_t bytes_written = cm_uart_write(dev, data, len, 0);
    lua_pushinteger(L, bytes_written);
    return 1;
}

// Lua: uart.read(id, max_len)
static int l_uart_read(lua_State *L)
{
    cm_uart_dev_e dev = (cm_uart_dev_e)luaL_checkinteger(L, 1);
    int max_len = luaL_checkinteger(L, 2);

    char* buffer = malloc(max_len);
    if (!buffer) return luaL_error(L, "failed to allocate read buffer");

    int32_t bytes_read = cm_uart_read(dev, buffer, max_len, 0);
    if (bytes_read >= 0) {
        lua_pushlstring(L, buffer, bytes_read);
    } else {
        lua_pushnil(L);
    }
    free(buffer);
    return 1;
}

// Lua: uart.on(id, "receive", callback)
static int l_uart_on(lua_State *L)
{
    cm_uart_dev_e dev = (cm_uart_dev_e)luaL_checkinteger(L, 1);
    const char *event_name = luaL_checkstring(L, 2);

    if (strcmp(event_name, "receive") != 0) {
        return luaL_argerror(L, 2, "only 'receive' event is supported");
    }

    if (g_uart_cbs[dev].on_receive_ref != LUA_NOREF) {
        luaL_unref(L, LUA_REGISTRYINDEX, g_uart_cbs[dev].on_receive_ref);
    }

    if (lua_isfunction(L, 3)) {
        lua_pushvalue(L, 3);
        g_uart_cbs[dev].on_receive_ref = luaL_ref(L, LUA_REGISTRYINDEX);
    } else {
        g_uart_cbs[dev].on_receive_ref = LUA_NOREF;
    }

    return 0;
}

// Lua: bytes_count = uart.readable(id)
static int l_uart_readable(lua_State *L)
{
    cm_uart_dev_e dev = (cm_uart_dev_e)luaL_checkinteger(L, 1);
    int32_t len = cm_uart_get_rxrb_data_len(dev);
    if (len < 0) {
        len = 0;
    }
    lua_pushinteger(L, len);
    return 1;
}

static const luaL_Reg uart_lib[] = {
    {"setup", l_uart_setup},
    {"write", l_uart_write},
    {"read", l_uart_read},
    {"on", l_uart_on},
    {"readable", l_uart_readable},
    {NULL, NULL}
};

int luaopen_uart(lua_State *L) {
    for (int i = 0; i < MAX_UARTS; i++) {
        g_uart_cbs[i].on_receive_ref = LUA_NOREF;
    }
    luaL_newlib(L, uart_lib);
    return 1;
}
